#!/usr/bin/env python3
# -*-coding:utf8-*-
# 通过速度控制机械臂末端位姿

from typing import (
    Optional,
)
import time
import math
from pynput import keyboard
from piper_sdk import *

def enable_fun(piper: C_PiperInterface_V2):
    '''
    使能机械臂并检测使能状态,尝试5s,如果使能超时则退出程序
    '''
    enable_flag = False
    # 设置超时时间（秒）
    timeout = 5
    # 记录进入循环前的时间
    start_time = time.time()
    elapsed_time_flag = False
    while not (enable_flag):
        elapsed_time = time.time() - start_time
        print("--------------------")
        enable_flag = piper.GetArmLowSpdInfoMsgs().motor_1.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_2.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_3.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_4.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_5.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_6.foc_status.driver_enable_status
        print("使能状态:", enable_flag)
        piper.EnableArm(7)
        piper.GripperCtrl(0, 1000, 0x01, 0)
        print("--------------------")
        # 检查是否超过超时时间
        if elapsed_time > timeout:
            print("超时....")
            elapsed_time_flag = True
            enable_flag = True
            break
        time.sleep(1)
        pass
    if(elapsed_time_flag):
        print("程序自动使能超时,退出程序")
        exit(0)

def print_instructions():
    """打印操作说明"""
    print("\n==== 机械臂末端速度控制 ====")
    print("使用键盘控制机械臂末端位姿:")
    print("W/S: 控制 X 轴正向/负向移动")
    print("A/D: 控制 Y 轴正向/负向移动")
    print("Q/E: 控制 Z 轴正向/负向移动")
    print("I/K: 控制 RX 轴正向/负向旋转")
    print("J/L: 控制 RY 轴正向/负向旋转")
    print("U/O: 控制 RZ 轴正向/负向旋转")
    print("1/2: 增加/减少速度")
    print("空格键: 停止移动")
    print("ESC: 退出程序")
    print("==========================\n")

if __name__ == "__main__":
    # 初始化机械臂接口
    piper = C_PiperInterface_V2("can0")
    piper.ConnectPort()
    piper.EnableArm(7)
    enable_fun(piper=piper)
    piper.GripperCtrl(0, 1000, 0x01, 0)
    
    # 设置转换因子，用于将物理单位转换为机械臂控制单位
    factor = 1000
    
    # 初始位置
    current_pose = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    
    # 获取当前末端位姿
    end_pose = piper.GetArmEndPoseMsgs()
    if end_pose:
        current_pose[0] = end_pose.end_pose.X_axis / factor
        current_pose[1] = end_pose.end_pose.Y_axis / factor
        current_pose[2] = end_pose.end_pose.Z_axis / factor
        current_pose[3] = end_pose.end_pose.RX_axis / factor
        current_pose[4] = end_pose.end_pose.RY_axis / factor
        current_pose[5] = end_pose.end_pose.RZ_axis / factor
    
    # 设置速度参数
    speed_factor = 0.5  # 初始速度因子
    move_step = 1.0     # 每次移动的基础步长 (mm)
    rotate_step = 1.0   # 每次旋转的基础步长 (度)
    
    # 设置控制模式为位置速度模式
    piper.MotionCtrl_2(0x01, 0x00, 50, 0x00)
    
    # 打印操作说明
    print_instructions()
    
    try:
    
        # 获取当前末端位姿
        end_pose = piper.GetArmEndPoseMsgs()
        if end_pose:
            print(f"当前位姿: X={end_pose.end_pose.X_axis/factor:.2f}mm, Y={end_pose.end_pose.Y_axis/factor:.2f}mm, Z={end_pose.end_pose.Z_axis/factor:.2f}mm, " +
                    f"RX={end_pose.end_pose.RX_axis/factor:.2f}°, RY={end_pose.end_pose.RY_axis/factor:.2f}°, RZ={end_pose.end_pose.RZ_axis/factor:.2f}°")
            print(f"当前速度因子: {speed_factor:.2f}")
            
            # 更新当前位姿
            current_pose[0] = end_pose.end_pose.X_axis / factor
            current_pose[1] = end_pose.end_pose.Y_axis / factor
            current_pose[2] = end_pose.end_pose.Z_axis / factor
            current_pose[3] = end_pose.end_pose.RX_axis / factor
            current_pose[4] = end_pose.end_pose.RY_axis / factor
            current_pose[5] = end_pose.end_pose.RZ_axis / factor
        
        # 使用一个简单的方法来检测按键状态
        pressed_keys = []
        
        def on_press(key):
            try:
                print(f"key.char: {key.char}")
                pressed_keys.append(key.char)
            except AttributeError:
                pressed_keys.append(key)
        
        def on_release(key):
            try:
                if key.char in pressed_keys:
                    pressed_keys.remove(key.char)
            except AttributeError:
                if key in pressed_keys:
                    pressed_keys.remove(key)
            if key == keyboard.Key.esc:
                return False
        
        # 创建监听器
        listener = keyboard.Listener(
            on_press=on_press,
            on_release=on_release)
        listener.start()
        
        # 主循环
        while True:
            # 根据键盘输入控制机械臂末端位姿
            if 'w' in pressed_keys:  # X轴正向
                current_pose[0] += move_step * speed_factor
            if 's' in pressed_keys:  # X轴负向
                current_pose[0] -= move_step * speed_factor
            if 'a' in pressed_keys:  # Y轴正向
                current_pose[1] += move_step * speed_factor
            if 'd' in pressed_keys:  # Y轴负向
                current_pose[1] -= move_step * speed_factor
            if 'q' in pressed_keys:  # Z轴正向
                current_pose[2] += move_step * speed_factor
            if 'e' in pressed_keys:  # Z轴负向
                current_pose[2] -= move_step * speed_factor
            if 'i' in pressed_keys:  # RX轴正向
                current_pose[3] += rotate_step * speed_factor
            if 'k' in pressed_keys:  # RX轴负向
                current_pose[3] -= rotate_step * speed_factor
            if 'j' in pressed_keys:  # RY轴正向
                current_pose[4] += rotate_step * speed_factor
            if 'l' in pressed_keys:  # RY轴负向
                current_pose[4] -= rotate_step * speed_factor
            if 'u' in pressed_keys:  # RZ轴正向
                current_pose[5] += rotate_step * speed_factor
            if 'o' in pressed_keys:  # RZ轴负向
                current_pose[5] -= rotate_step * speed_factor
            if '1' in pressed_keys:  # 增加速度
                speed_factor = min(speed_factor + 0.1, 5.0)
                pressed_keys.remove('1')
                time.sleep(0.2)  # 防止速度变化太快
            if '2' in pressed_keys:  # 减小速度
                speed_factor = max(speed_factor - 0.1, 0.1)
                pressed_keys.remove('2')
                time.sleep(0.2)  # 防止速度变化太快
            if keyboard.Key.space in pressed_keys:  # 停止
                print("停止移动")
                # 不更新位置，保持当前位置
                pass
            if keyboard.Key.esc in pressed_keys:  # 退出
                print("程序退出")
                break
        
            # 将物理单位转换为机械臂控制单位
            X = round(current_pose[0] * factor)
            Y = round(current_pose[1] * factor)
            Z = round(current_pose[2] * factor)
            RX = round(current_pose[3] * factor)
            RY = round(current_pose[4] * factor)
            RZ = round(current_pose[5] * factor)
            
            # 设置机械臂末端位姿
            piper.MotionCtrl_2(0x01, 0x00, int(50 * speed_factor), 0x00)  # 根据速度因子调整速度百分比
            piper.EndPoseCtrl(X, Y, Z, RX, RY, RZ)
            
            # 短暂延时，避免过于频繁的控制命令
            time.sleep(0.05)
            
    except KeyboardInterrupt:
        print("程序被用户中断")
    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        # 程序结束前，确保机械臂停止
        piper.MotionCtrl_2(0, 0, 0, 0x00)
        print("程序已安全退出")
